from _ast import AST, ImportFrom
from ast import NodeVisitor
from typing import Any
import ast
import astunparse

class ClassVisitor(NodeVisitor):
    def __init__(self):
        self.result = {}

    def get_args(self, node):
        args = node.args
        arg_names = []
        defaults = args.defaults
        for arg in args.args:
            arg_names += [arg.arg]
        return arg_names, len(defaults)

    def visit_FunctionDef(self, node):
        kw_names = self.get_args(node)
        self.result[node.name] = kw_names
        return node

class FunVisitor(NodeVisitor):
    def __init__(self):
        self.result = {}
        self.pair = {}

    def get_args(self, node):
        args = node.args
        arg_names = []
        defaults = args.defaults
        for arg in args.args:
            arg_names += [arg.arg]
        return arg_names, len(defaults)

    def visit_FunctionDef(self, node):
        kw_names = self.get_args(node)
        self.result[node.name] = kw_names
        return node

    def visit_ClassDef(self, node):
        visitor = ClassVisitor()
        visitor.visit(node)
        # print(node.name, node.bases)
        if len(node.bases) > 0: # bases 为一个由显式指明的基类节点组成的列表
            self.pair[node.name] = []
            for base in node.bases:
                if isinstance(base, ast.Attribute):
                    self.pair[node.name].append(self._deal_Attribute(base))
                elif isinstance(base, ast.Name):
                    self.pair[node.name].append(base.id)
        # else:
        #     self.pair[node.name] = None
        self.result[node.name] = visitor.result
        return node
    
    def _deal_Attribute(self, node):
        """ Attribute(
                value=Name(id='snake', ctx=Load()),
                attr='colour',
                ctx=Load()))
        """
        if isinstance(node, ast.Attribute):
            prefix = self._deal_Attribute(node.value)
            return prefix + '.' + node.attr 
        elif isinstance(node, ast.Name):
            return node.id


class ImportVisitor(NodeVisitor):
    def __init__(self, dp_dict):
        self.import_dict = {} # {name:asname}
        self.importfrom_dict = {} # {level: {prefix_path: [[name, asname]]}
        self.dp_dict = dp_dict

    def visit_Import(self, node):
        # print(astunparse.dump(node))
        # print(node.names)
        for alias_node in node.names:
            # print(alias_node.name, alias_node.asname)
            self.import_dict[alias_node.name] = alias_node.asname

        return node

    def visit_ImportFrom(self, node):
        # print(astunparse.dump(node))
        if not node.level in self.importfrom_dict:
            self.importfrom_dict[node.level] = {}

        if node.module is not None:
            for dp in self.dp_dict:
                # print(type(node.module), node.module, dp)
                if node.module.startswith(dp):
                    node.module = node.module.replace(dp, dp + "-" + self.dp_dict[dp], 1)

        if not node.module in self.importfrom_dict[node.level]:
            self.importfrom_dict[node.level][node.module] = []
        for alias_node in node.names:
            self.importfrom_dict[node.level][node.module].append(([alias_node.name, alias_node.asname]))
        
        return node
